#ifndef SKIPGRID_H
#define SKIPGRID_H

/**
 * \bibitem{de2017skimap}
 * D. De Gregorio and L. Di Stefano, ``Skimap: An Efficient Mapping Framework for Robot Navigation,'' In \textit{Proceedings of the 2017 IEEE International Conference on Robotics and Automation (ICRA)}, pp. 2569-2576, 2017. IEEE.
 */

#include "key.h"
#include "cell.h"
#include "table.h"
#include "skiplist.h"
#include <malloc.h>
#include <omp.h>

namespace RealScan
{

// SkipGrid implementation of `Table`
template <typename T, unsigned L = 16, unsigned THREADS_NUM = 12>
class SkipGrid : public Table<T>
{
public:
    // Y branch type
    using YBranch = SkipList<T, L>;
    // X branch type
    using XBranch = SkipList<YBranch*, L>;

    // Constructor
    SkipGrid() 
    {
        // Initialize the root node
        m_root = new XBranch();
    }

    // Destructor
    ~SkipGrid()
    {
        // Release Y branch nodes
        SkipNode<YBranch*>* xNode = m_root->head()->forward[0];

        while (xNode != m_root->tail())
        {
            if (xNode->value)
            {
                delete xNode->value;
            }
            xNode = xNode->forward[0];
        }

        // Release X branch (root)
        delete m_root;

        malloc_trim(0);
    }

    ////// `Table` interface implementation //////

    // Search
    std::optional<T> search(const Key& key) const override
    {
        // Search in the X branch
        SkipNode<YBranch*>* xNode = m_root->search(key.x);

        // If the X node exists and the Y branch is present
        if (xNode && xNode->value)
        {
            // Search in the Y branch
            SkipNode<T>* yNode = xNode->value->search(key.y);

            return (yNode ? std::optional<T>(yNode->value) : std::nullopt);
        }
        return std::nullopt;
    }

    // Insert
    bool insert(const Key& key, const T& value) override
    {
        // Search for the X branch
        SkipNode<YBranch*>* xNode = m_root->search(key.x);

        // If the X branch doesn't exist, create it
        if (!xNode) {
            xNode = m_root->insert(key.x, new YBranch);
        }

        // If insertion is successful, update the Y branch
        if (xNode) {
            xNode->value->insert(key.y, value);
        }
        return true;
    }

    // Remove
    void remove(const Key& key) override
    {
        // Search for the X branch
        SkipNode<YBranch*>* xNode = m_root->search(key.x);

        // If the branch exists, attempt to delete
        if (xNode && xNode->value)
        {
            xNode->value->remove(key.y);

            // If no elements are left in the Y branch, consider removing the entire branch
            if (xNode->value->size() == 0)
            {
                delete xNode->value;
                m_root->remove(key.x);
            }
        }
    }

    // Vectorize
    void toVector(std::vector<Cell<T> >& vec, const Key& minXY, const Key& maxXY) const override
    {
        // Clear the vector
        vec.clear();

        // Get the list of X nodes
        std::vector<SkipNode<YBranch*>*> xNodes;
        m_root->toVector(xNodes, minXY.x, maxXY.x);

        // Process the Y branches in parallel
        #pragma omp parallel num_threads(THREADS_NUM)
        {
            std::vector<Cell<T>> localCells;

            #pragma omp for nowait
            for (std::size_t i = 0; i < xNodes.size(); ++i)
            {
                if (xNodes[i] && xNodes[i]->value)
                {
                    // Get the list of Y nodes
                    std::vector<SkipNode<T>*> yNodes;
                    xNodes[i]->value->toVector(yNodes, minXY.y, maxXY.y);

                    for (std::size_t j = 0; j < yNodes.size(); ++j)
                    {
                        localCells.emplace_back(Cell{Key{xNodes[i]->key, yNodes[j]->key}, yNodes[j]->value});
                    }
                }
            }

            // Merge local grids into the main vector
            #pragma omp critical
            {
                vec.insert(vec.end(), localCells.begin(), localCells.end());
            }
        }
    }

    // cells number
    std::size_t size() const override
    {
        std::size_t size_ = 0;

        // Traverse X branches
        SkipNode<YBranch*>* xNode = m_root->head()->forward[0];

        // Sum the sizes
        while (xNode != m_root->tail())
        {
            if (xNode->value)
            {
                size_ += xNode->value->size();
            }
            xNode = xNode->forward[0];
        }

        return size_;
    }

    // memory used, const there
    std::size_t space() const override
    {
        std::size_t space_ = 0;

        // Traverse X branches
        SkipNode<YBranch*>* xNode = m_root->head()->forward[0];

        // Sum the space usage for Y branches
        while (xNode != m_root->tail())
        {
            if (xNode->value)
            {
                space_ += xNode->value->space();
            }
            xNode = xNode->forward[0];
        }

        // Add space used by the X branches
        space_ += m_root->space();
        return space_;
    }

private:
    // Root node
    XBranch* m_root = nullptr;
};

} /* RealScan */

#endif /* SKIPGRID_H */
